# -*- coding: utf-8 -*-
#
# @Author: CPS
# @email: 373704015@qq.com
# @Date: 2024-04-17 15:43:16.776498
# @Last Modified by: CPS
# @Last Modified time: 2024-04-17 15:43:16.776498
# @file_path "Z:\work\2024\项目\中大模型数据提取处理"
# @Filename "main.py"
# @Description: .g13文件解析功能模块
#
import os, sys

sys.path.append("..")

from os import path
import pandas as pd
from Types import RiverDataType, SectionDataType


def get_river_data(data: list[str]) -> dict:
    """
    @Description 处理从g13（hex模型）导出的文本文件，从中提取河段和断面数据

    - param data :{list[str]} 通过`f.readlines()`后得到的列表文件

    @returns `{ dict}`

    """
    river_start_symbol = "River Reach="
    xy_start_symbol = "Reach XY="
    # xy_end_symbol = "Rch Text X Y="

    section_start_symbol = "Type RM Length L Ch R ="  # 这行标记断面的边距
    section_coords_start_symbol = "#Sta/Elev="  # 这里记录当前断面有多少个高程坐标 x,z
    section_coords_end_symbol = "#Mann="  # 同时也是曼宁系数的位置（下一行）

    currt_river_name = ""
    currt_section_name = ""
    currt_section_info = None

    next_action = "find_river"
    # 'find_river' 当前没在河流中
    # 'record_river_name' 记录发现河道名称，河段
    # 'record_river_xy_len' 记录河道xy长度
    # 'record_river_xy' 开始记录xy信息

    river_info = {}  # {"河流名称":{河流数据}}

    for currt_line in data:
        if "Stream Node=" in currt_line:
            break

        # 过滤注释
        if "'" in currt_line:
            currt_line = currt_line.split("'")[0]

        # 发现河道
        if river_start_symbol in currt_line:
            # 每次查找到新河流，就临时变量清空
            currt_river_name = ""
            currt_section_name = ""
            currt_section_info = None

            tmp_list = currt_line.split(",")
            river_name = tmp_list[0].strip().split("=")[1]  # 获取河流名
            river_part = tmp_list[1].strip()  # 获取河段

            # 获取河流信息
            if not river_name in river_info:
                currt_river_name = f"{river_name}_{river_part}"
                river_info[currt_river_name] = {
                    "river_name": river_name,
                    "river_part": river_part,
                    "river_xy": [],
                }

        elif xy_start_symbol in currt_line:
            # 记录xy长度
            tmp_list = currt_line.split("= ")
            river_info[currt_river_name]["river_xy_len"] = int(tmp_list[1])

            next_action = "record_river_xy"
        elif next_action == "record_river_xy":
            # 记录xy数据
            xy_list = currt_line.split()
            if len(xy_list) >= 2:
                river_info[currt_river_name]["river_xy"] += xy_list

            if len(river_info[currt_river_name]["river_xy"]) == river_info[currt_river_name]["river_xy_len"] * 2:
                # 已经提取足够的坐标
                next_action = "record_river_xy_end"

        elif next_action == "record_river_xy_end":
            # 修正xy顺序为[(x1,y1), (x2,y2), ...]
            xy_list = river_info[currt_river_name]["river_xy"]
            river_info[currt_river_name]["river_xy"] = list(zip(xy_list[::2], xy_list[1::2]))
            river_info[currt_river_name]["river_xy_len"] = len(river_info[currt_river_name]["river_xy"])
            river_info[currt_river_name]["section_info"] = {}
            # 验证数量
            if river_info[currt_river_name]["river_xy_len"] != len(river_info[currt_river_name]["river_xy"]):
                print(f"【警告】 当前河流{currt_river_name}xy长度不匹配")

            next_action = "find_section"

        elif next_action == "find_section" and section_start_symbol in currt_line:
            currt_line_list = [e.strip() for e in currt_line.split(",")]

            # 断面数据
            sid = currt_line_list[-4]
            distance = currt_line_list[-3:]
            currt_section_name = f"{currt_river_name}_section:{sid}"
            river_name, river_part = currt_river_name.split("_")

            section_info = {
                "river_name": river_name,
                "river_part": river_part,
                "section_name": currt_section_name,
                "sid": sid,
                "next_section_distance_left": float(distance[0]),
                "next_section_distance_center": float(distance[1]),
                "next_section_distance_right": float(distance[2]),
                "xz_list": [],
                "is_warnning": False,
            }

            river_info[currt_river_name]["section_info"][currt_section_name] = section_info
            next_action = "record_section_xz_len"

        elif next_action == "record_section_xz_len" and section_coords_start_symbol in currt_line:
            currt_line_list = [e.strip() for e in currt_line.split("=")]

            currt_section_info = river_info[currt_river_name]["section_info"][currt_section_name]
            currt_section_info["xz_len"] = int(currt_line_list[-1])

            next_action = "record_section_xz"

        elif next_action == "record_section_xz":
            currt_line_list = [e.strip() for e in currt_line.split()]
            currt_section_info = river_info[currt_river_name]["section_info"][currt_section_name]

            if len(currt_line_list) == 1:
                # print(f"【{currt_section_name}】: xy数据异常，不进行记录")
                currt_section_info["xz_list"] = [("0", "0")] * currt_section_info["xz_len"]
                currt_section_info["is_warnning"] = True
                next_action = "record_section_xz_end"
                continue

            currt_section_info["xz_list"] += currt_line_list

            if len(currt_section_info["xz_list"]) == currt_section_info["xz_len"] * 2:
                currt_section_info["xz_list"] = list(
                    zip(currt_section_info["xz_list"][::2], currt_section_info["xz_list"][1::2])
                )
                next_action = "record_section_xz_end"

        elif section_coords_end_symbol in currt_line:
            # 曼宁系数在这行的下一行
            next_action = "record_section_Mann"

        elif next_action == "record_section_Mann":
            # 记录曼宁
            currt_line_list = [e.strip() for e in currt_line.split()]
            river_info[currt_river_name]["section_info"][currt_section_name]["Mann"] = currt_line_list[1]

            next_action = "find_section"
        else:
            continue

    return river_info


def data_check(data: dict):
    """
    @Description 检查数据

    - param data    :{dict} 目标数据
    - param count=0 :{int}  检查个数

    """
    i = 1
    for river_key, river_data in data.items():
        name, key = river_key.split("_")

        # 检查断面数量是否符合
        if river_data["river_xy_len"] > len(river_data["section_info"].keys()):
            print("--- --- --- --- --- --- --- --- --- --- --- --- --- --- ---")
            print(f"{i} warning !")
            print(f"河流:{name} 河段：{key} 缺少数据")
            print("断面数量【实际】: ", river_data["river_xy_len"])
            print("断面数量【当前】: ", len(river_data["section_info"].keys()))
            print(f"搜索代码: {river_start_symbol}{name}\\s+,{key}")
            # print(river_data)
            for e in river_data["section_info"].keys():
                print(e)
            print("--- --- --- --- --- --- --- --- --- --- --- --- --- --- ---")

        elif river_data["river_xy_len"] < len(river_data["section_info"].keys()):
            print("--- --- --- --- --- --- --- --- --- --- --- --- --- --- ---")
            print(f"{i} warning !")
            print(f"河流:{name} 河段：{key} 过多数据")
            print("断面数量【实际】: ", river_data["river_xy_len"])
            print("断面数量【当前】: ", len(river_data["section_info"].keys()))
            print(f"搜索代码: {river_start_symbol}{name}\\s+,{key}")
            print("--- --- --- --- --- --- --- --- --- --- --- --- --- --- ---")

        else:
            print(f"{i} done !")

        # if i == 100:
        #     break

        i += 1

    print(len(res.keys()))


def data_format(data: dict) -> list[RiverDataType]:
    """
    @Description 格式化数据，将数据格式化成`RiverDataType`格式

    - param data :{dict} {description}

    @returns `{ list[RiverDataType]}` {description}

    """
    result = []
    for river_data in data.values():
        section_list = [
            SectionDataType(**river_data["section_info"][section_key])
            for section_key in river_data["section_info"].keys()
        ]

        river_data_obj = RiverDataType(
            **{
                "river_name": river_data["river_name"],
                "river_part": river_data["river_part"],
                "river_xy_len": river_data["river_xy_len"],
                "river_xy": river_data["river_xy"],
                "section_list": section_list,
            }
        )

        result.append(river_data_obj)

    return result


def river_data_to_excel(data: list[RiverDataType], output: str = "river_result.xlsx"):
    """
    @Description 将格式化后的【断面】数据，导出成xlsx格式

    - param data   :{list[RiverDataType]} {description}
    - param output :{str}                 {description}

    """

    river_name_col = []
    river_part_col = []
    X_col = []
    Y_col = []

    row = 0
    for each_data in data:
        # 获取行数长度，重来填充
        cols_len = len(each_data.river_xy)

        river_name_col += [each_data.river_name] * cols_len
        river_part_col += [each_data.river_part] * cols_len
        X_col += [float(each_xy[0]) for each_xy in each_data.river_xy]
        Y_col += [float(each_xy[1]) for each_xy in each_data.river_xy]

        row += 1
        # if row == max_col:
        #     break

    df = pd.DataFrame(
        {
            "河道名称": river_name_col,
            "河段名称": river_part_col,
            "X": X_col,
            "Y": Y_col,
        }
    )
    df.to_excel(output, index=False)


def section_data_to_excel(data: list[RiverDataType], output: str = "river_section_result.xlsx"):
    """
    @Description 将格式化后的【河道】数据，导出成xlsx格式

    - param data   :{list[RiverDataType]} 已格式化的数据
    - param output :{str}                 {description}

    @returns `{}` {description}

    """
    max_row = 300

    river_name_col = []
    river_part_col = []

    sid_col = []
    next_section_distance_left_col = []
    next_section_distance_center_col = []
    next_section_distance_right_col = []
    Mann_col = []

    X_col = []
    Z_col = []

    is_warnning_col = []

    row = 0
    for each_data in data:
        for section_data in each_data.section_list:
            # 获取行数长度，重来填充
            cols_len = len(section_data.xz_list)

            if section_data.is_warnning == True:
                print(f"发现导*行: {each_data.river_name}_{each_data.river_part}_{section_data.sid}")
                is_warnning_col += ["*"] * cols_len
            else:
                is_warnning_col += [""] * cols_len

            river_name_col += [each_data.river_name] * cols_len
            river_part_col += [each_data.river_part] * cols_len

            sid_col += [section_data.sid] * cols_len
            next_section_distance_left_col += [section_data.next_section_distance_left] * cols_len

            next_section_distance_center_col += [section_data.next_section_distance_center] * cols_len
            next_section_distance_right_col += [section_data.next_section_distance_right] * cols_len
            Mann_col += [section_data.Mann] * cols_len

            X_col += [float(each_xz[0]) for each_xz in section_data.xz_list]
            Z_col += [float(each_xz[1]) for each_xz in section_data.xz_list]

            cols_len = 0
        row += 1

        if row >= max_row:
            break

    try:
        df = pd.DataFrame(
            {
                "河道名称": river_name_col,
                "河段名称": river_part_col,
                "断面号": sid_col,
                "左间距": next_section_distance_left_col,
                "中间距": next_section_distance_center_col,
                "右间距": next_section_distance_right_col,
                "糙率": Mann_col,
                "X": X_col,
                "Z": Z_col,
                "格式错误": is_warnning_col,
            }
        )
        df.to_excel(output, index=False)
    except Exception as e:
        print(e)
        print(
            {
                "河道名称": len(river_name_col),
                "河段名称": len(river_part_col),
                "断面号": len(sid_col),
                "左间距": len(next_section_distance_left_col),
                # "中间距": next_section_distance_center_col,
                # "右间距": next_section_distance_right_col,
                # "糙率": Mann_col,
                # "X": X_col,
                # "Z": Z_col,
            }
        )

    print(f"处理成功，{row}行数据")


if __name__ == "__main__":
    input_data = path.abspath("./data.g13")

    print("input_data: ", input_data)
    river_start_symbol = "River Reach="

    with open(input_data, "r", encoding="utf-8") as f:
        full_data = f.readlines()
        res = get_river_data(full_data)

    # data_check(res)
    formated_data = data_format(res)

    # 导出河道信息
    # river_data_to_excel(formated_data)

    # 导出断面信息
    section_data_to_excel(formated_data, "断面数据（错误标星）.xlsx")
